#include <mcp/auth/oauth2.h>
#include <mcp/error.h>
#include <mcp/json/value.h>

#include <iostream>
#include <sstream>
#include <chrono>
#include <random>
#include <memory>
#include <unordered_map>
#include <mutex>
#include <curl/curl.h>
#include <openssl/evp.h>
#include <openssl/bio.h>
#include <openssl/buffer.h>

namespace mcp {
namespace auth {

// Helper functions for OAuth2
namespace {
    static std::string generate_random_string(size_t length) {
        static const std::string chars =
            "ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789-._~";

        std::random_device rd;
        std::mt19937 rng(rd());
        std::uniform_int_distribution<size_t> dist(0, chars.length() - 1);

        std::string result;
        result.reserve(length);
        for (size_t i = 0; i < length; ++i) {
            result += chars[dist(rng)];
        }
        return result;
    }

    static std::string base64url_encode(const std::string& data) {
        BIO *bio, *b64;
        BUF_MEM *bufferPtr;

        b64 = BIO_new(BIO_f_base64());
        bio = BIO_new(BIO_s_mem());
        bio = BIO_push(b64, bio);

        BIO_set_flags(bio, BIO_FLAGS_BASE64_NO_NL);
        BIO_write(bio, data.c_str(), static_cast<int>(data.length()));
        BIO_flush(bio);
        BIO_get_mem_ptr(bio, &bufferPtr);

        std::string result(bufferPtr->data, bufferPtr->length);
        BIO_free_all(bio);

        // Replace + with - and / with _ and remove padding
        std::replace(result.begin(), result.end(), '+', '-');
        std::replace(result.begin(), result.end(), '/', '_');
        result.erase(std::remove(result.begin(), result.end(), '='), result.end());

        return result;
    }

    static std::string sha256(const std::string& data) {
        unsigned char hash[EVP_MAX_MD_SIZE];
        unsigned int hash_len;

        EVP_MD_CTX* ctx = EVP_MD_CTX_new();
        EVP_DigestInit_ex(ctx, EVP_sha256(), nullptr);
        EVP_DigestUpdate(ctx, data.c_str(), data.length());
        EVP_DigestFinal_ex(ctx, hash, &hash_len);
        EVP_MD_CTX_free(ctx);

        return std::string(reinterpret_cast<char*>(hash), hash_len);
    }

    // HTTP request helper
    struct HttpResponse {
        int status_code = 0;
        std::string body;
        std::unordered_map<std::string, std::string> headers;
    };

    static std::optional<HttpResponse> http_request(
        const std::string& url,
        const std::string& method = "GET",
        const std::unordered_map<std::string, std::string>& headers = {},
        const std::string& body = "") {

        CURL* curl = curl_easy_init();
        if (!curl) {
            return std::nullopt;
        }

        std::string response_body;
        std::unordered_map<std::string, std::string> response_headers;

        // Set URL
        curl_easy_setopt(curl, CURLOPT_URL, url.c_str());

        // Set method
        if (method == "POST") {
            curl_easy_setopt(curl, CURLOPT_POST, 1L);
            if (!body.empty()) {
                curl_easy_setopt(curl, CURLOPT_POSTFIELDS, body.c_str());
            }
        }

        // Set headers
        curl_slist* header_list = nullptr;
        for (const auto& [key, value] : headers) {
            std::string header_line = key + ": " + value;
            header_list = curl_slist_append(header_list, header_line.c_str());
        }
        if (header_list) {
            curl_easy_setopt(curl, CURLOPT_HTTPHEADER, header_list);
        }

        // Set response callbacks
        curl_easy_setopt(curl, CURLOPT_WRITEFUNCTION, +[](char* ptr, size_t size, size_t nmemb, void* userdata) -> size_t {
            std::string* response = static_cast<std::string*>(userdata);
            response->append(ptr, size * nmemb);
            return size * nmemb;
        });
        curl_easy_setopt(curl, CURLOPT_WRITEDATA, &response_body);

        curl_easy_setopt(curl, CURLOPT_HEADERFUNCTION, +[](char* buffer, size_t size, size_t nitems, void* userdata) -> size_t {
            auto* headers = static_cast<std::unordered_map<std::string, std::string>*>(userdata);
            std::string header(buffer, size * nitems);

            size_t colon_pos = header.find(':');
            if (colon_pos != std::string::npos) {
                std::string key = header.substr(0, colon_pos);
                std::string value = header.substr(colon_pos + 1);

                // Trim whitespace
                while (!value.empty() && std::isspace(value.front())) value.erase(0, 1);
                while (!value.empty() && std::isspace(value.back())) value.pop_back();

                (*headers)[key] = value;
            }

            return size * nitems;
        });
        curl_easy_setopt(curl, CURLOPT_HEADERDATA, &response_headers);

        // Set timeout
        curl_easy_setopt(curl, CURLOPT_TIMEOUT, 30L);

        // Perform request
        CURLcode res = curl_easy_perform(curl);

        HttpResponse response;
        if (res == CURLE_OK) {
            curl_easy_getinfo(curl, CURLINFO_RESPONSE_CODE, &response.status_code);
            response.body = std::move(response_body);
            response.headers = std::move(response_headers);
        }

        // Cleanup
        if (header_list) {
            curl_slist_free_all(header_list);
        }
        curl_easy_cleanup(curl);

        if (res != CURLE_OK) {
            return std::nullopt;
        }

        return response;
    }

    static std::unordered_map<std::string, std::string> parse_form_encoded(const std::string& data) {
        std::unordered_map<std::string, std::string> result;

        size_t start = 0;
        while (start < data.length()) {
            size_t amp_pos = data.find('&', start);
            if (amp_pos == std::string::npos) {
                amp_pos = data.length();
            }

            size_t equal_pos = data.find('=', start);
            if (equal_pos != std::string::npos && equal_pos < amp_pos) {
                std::string key = data.substr(start, equal_pos - start);
                std::string value = data.substr(equal_pos + 1, amp_pos - equal_pos - 1);
                result[key] = value;
            }

            start = amp_pos + 1;
        }

        return result;
    }
}

// OAuth2Authenticator implementation
OAuth2Authenticator::OAuth2Authenticator(const OAuth2Config& config) : config_(config) {}

Result<AuthResult> OAuth2Authenticator::authenticate(const Credentials& credentials) {
    if (config_.use_pkce) {
        // Authorization Code Flow with PKCE
        auto [verifier, challenge] = generate_pkce_codes();
        std::string state = generate_random_string(32);

        std::string auth_url = get_authorization_url(state, verifier);

        // In a real implementation, this would redirect the user to auth_url
        // and wait for callback with authorization code
        // For now, we'll simulate this with client credentials flow
        return Error::unsupported_operation("Authorization code flow requires user interaction");
    } else {
        // Client Credentials Flow
        std::string post_data =
            "grant_type=client_credentials" +
            "&client_id=" + curl_easy_escape(nullptr, config_.client_id.c_str(), 0) +
            "&client_secret=" + curl_easy_escape(nullptr, config_.client_secret.c_str(), 0);

        if (!config_.scopes.empty()) {
            std::string scope_str;
            for (const auto& scope : config_.scopes) {
                if (!scope_str.empty()) scope_str += " ";
                scope_str += scope;
            }
            post_data += "&scope=" + curl_easy_escape(nullptr, scope_str.c_str(), 0);
        }

        auto response = http_request(config_.token_url, "POST", {{"Content-Type", "application/x-www-form-urlencoded"}}, post_data);

        if (!response || response->status_code != 200) {
            return Error::authentication_error("Failed to obtain OAuth2 token: " +
                (response ? std::to_string(response->status_code) : "network error"));
        }

        auto params = parse_form_encoded(response->body);

        AuthResult result;
        result.success = true;
        result.access_token = params["access_token"];
        result.token_type = params["token_type"];
        result.expires_in = std::stoi(params["expires_in"]);

        if (params.count("refresh_token")) {
            result.refresh_token = params["refresh_token"];
        }

        return result;
    }
}

Result<AuthResult> OAuth2Authenticator::refresh(const std::string& refresh_token) {
    if (refresh_token.empty()) {
        return Error::validation_error("Refresh token is required");
    }

    std::string post_data =
        "grant_type=refresh_token" +
        "&refresh_token=" + curl_easy_escape(nullptr, refresh_token.c_str(), 0) +
        "&client_id=" + curl_easy_escape(nullptr, config_.client_id.c_str(), 0);

    if (!config_.client_secret.empty()) {
        post_data += "&client_secret=" + curl_easy_escape(nullptr, config_.client_secret.c_str(), 0);
    }

    auto response = http_request(config_.token_url, "POST", {{"Content-Type", "application/x-www-form-urlencoded"}}, post_data);

    if (!response || response->status_code != 200) {
        return Error::authentication_error("Failed to refresh OAuth2 token");
    }

    auto params = parse_form_encoded(response->body);

    AuthResult result;
    result.success = true;
    result.access_token = params["access_token"];
    result.token_type = params["token_type"];
    result.expires_in = std::stoi(params["expires_in"]);

    if (params.count("refresh_token")) {
        result.refresh_token = params["refresh_token"];
    }

    return result;
}

bool OAuth2Authenticator::validate_token(const std::string& access_token) {
    // Basic validation - in real implementation, would validate JWT or check with auth server
    return !access_token.empty() && access_token.length() > 10;
}

std::string OAuth2Authenticator::get_scheme() const {
    return "Bearer";
}

std::vector<AuthChallenge> OAuth2Authenticator::get_challenges() const {
    return {AuthChallenge{
        .scheme = "Bearer",
        .realm = "OAuth2",
        .parameters = {
            {"authorization_url", config_.auth_url},
            {"token_url", config_.token_url}
        }
    }};
}

bool OAuth2Authenticator::supports_refresh() const {
    return true;
}

std::string OAuth2Authenticator::get_authorization_url(const std::string& state,
                                                       const std::string& code_verifier) const {
    std::string url = config_.auth_url + "?response_type=code" +
                      "&client_id=" + curl_easy_escape(nullptr, config_.client_id.c_str(), 0) +
                      "&redirect_uri=" + curl_easy_escape(nullptr, config_.redirect_uri.c_str(), 0) +
                      "&state=" + curl_easy_escape(nullptr, state.c_str(), 0);

    if (config_.use_pkce) {
        url += "&code_challenge=" + curl_easy_escape(nullptr,
                generate_code_challenge(code_verifier).c_str(), 0) +
               "&code_challenge_method=" + config_.code_challenge_method;
    }

    if (!config_.scopes.empty()) {
        std::string scope_str;
        for (const auto& scope : config_.scopes) {
            if (!scope_str.empty()) scope_str += " ";
            scope_str += scope;
        }
        url += "&scope=" + curl_easy_escape(nullptr, scope_str.c_str(), 0);
    }

    return url;
}

std::expected<OAuth2TokenResponse, OAuth2Error> OAuth2Authenticator::exchange_code(
    const std::string& code,
    const std::string& code_verifier,
    const std::string& redirect_uri) {

    std::string post_data =
        "grant_type=authorization_code" +
        "&code=" + curl_easy_escape(nullptr, code.c_str(), 0) +
        "&redirect_uri=" + curl_easy_escape(nullptr, redirect_uri.c_str(), 0) +
        "&client_id=" + curl_easy_escape(nullptr, config_.client_id.c_str(), 0);

    if (config_.use_pkce) {
        post_data += "&code_verifier=" + curl_easy_escape(nullptr, code_verifier.c_str(), 0);
    }

    if (!config_.client_secret.empty()) {
        post_data += "&client_secret=" + curl_easy_escape(nullptr, config_.client_secret.c_str(), 0);
    }

    auto response = http_request(config_.token_url, "POST", {{"Content-Type", "application/x-www-form-urlencoded"}}, post_data);

    if (!response) {
        return OAuth2Error{"network_error", "Failed to connect to token endpoint"};
    }

    if (response->status_code != 200) {
        auto params = parse_form_encoded(response->body);
        OAuth2Error error;
        error.error = params.count("error") ? params["error"] : "unknown";
        if (params.count("error_description")) {
            error.error_description = params["error_description"];
        }
        return error;
    }

    auto params = parse_form_encoded(response->body);
    OAuth2TokenResponse token_response;
    token_response.access_token = params["access_token"];
    token_response.token_type = params["token_type"];
    token_response.expires_in = std::stoi(params["expires_in"]);

    if (params.count("refresh_token")) {
        token_response.refresh_token = params["refresh_token"];
    }
    if (params.count("scope")) {
        token_response.scope = params["scope"];
    }

    return token_response;
}

std::expected<OAuth2TokenResponse, OAuth2Error> OAuth2Authenticator::refresh_token(
    const std::string& refresh_token) {

    std::string post_data =
        "grant_type=refresh_token" +
        "&refresh_token=" + curl_easy_escape(nullptr, refresh_token.c_str(), 0) +
        "&client_id=" + curl_easy_escape(nullptr, config_.client_id.c_str(), 0);

    if (!config_.client_secret.empty()) {
        post_data += "&client_secret=" + curl_easy_escape(nullptr, config_.client_secret.c_str(), 0);
    }

    auto response = http_request(config_.token_url, "POST", {{"Content-Type", "application/x-www-form-urlencoded"}}, post_data);

    if (!response || response->status_code != 200) {
        OAuth2Error error;
        error.error = "invalid_grant";
        error.error_description = "Failed to refresh token";
        return error;
    }

    auto params = parse_form_encoded(response->body);
    OAuth2TokenResponse token_response;
    token_response.access_token = params["access_token"];
    token_response.token_type = params["token_type"];
    token_response.expires_in = std::stoi(params["expires_in"]);

    if (params.count("refresh_token")) {
        token_response.refresh_token = params["refresh_token"];
    }
    if (params.count("scope")) {
        token_response.scope = params["scope"];
    }

    return token_response;
}

std::pair<std::string, std::string> OAuth2Authenticator::generate_pkce_codes() {
    std::string verifier = generate_code_verifier();
    std::string challenge = generate_code_challenge(verifier);
    return {verifier, challenge};
}

const OAuth2Config& OAuth2Authenticator::get_config() const {
    return config_;
}

void OAuth2Authenticator::update_config(const OAuth2Config& config) {
    config_ = config;
}

std::string OAuth2Authenticator::generate_code_verifier() const {
    return generate_random_string(128);
}

std::string OAuth2Authenticator::generate_code_challenge(const std::string& verifier) const {
    std::string hash = sha256(verifier);
    return base64url_encode(hash);
}

std::string OAuth2Authenticator::base64url_encode(const std::string& data) const {
    return ::mcp::auth::base64url_encode(data);
}

// OAuth2ClientCredentialsAuthenticator implementation
OAuth2ClientCredentialsAuthenticator::OAuth2ClientCredentialsAuthenticator(const OAuth2Config& config)
    : config_(config) {}

Result<AuthResult> OAuth2ClientCredentialsAuthenticator::authenticate(const Credentials& credentials) {
    std::string post_data =
        "grant_type=client_credentials" +
        "&client_id=" + curl_easy_escape(nullptr, config_.client_id.c_str(), 0) +
        "&client_secret=" + curl_easy_escape(nullptr, config_.client_secret.c_str(), 0);

    if (!config_.scopes.empty()) {
        std::string scope_str;
        for (const auto& scope : config_.scopes) {
            if (!scope_str.empty()) scope_str += " ";
            scope_str += scope;
        }
        post_data += "&scope=" + curl_easy_escape(nullptr, scope_str.c_str(), 0);
    }

    auto response = http_request(config_.token_url, "POST", {{"Content-Type", "application/x-www-form-urlencoded"}}, post_data);

    if (!response || response->status_code != 200) {
        return Error::authentication_error("Failed to obtain client credentials token");
    }

    auto params = parse_form_encoded(response->body);

    AuthResult result;
    result.success = true;
    result.access_token = params["access_token"];
    result.token_type = params["token_type"];
    result.expires_in = std::stoi(params["expires_in"]);

    return result;
}

Result<AuthResult> OAuth2ClientCredentialsAuthenticator::refresh(const std::string& refresh_token) {
    return authenticate(Credentials{}); // Client credentials flow doesn't use refresh tokens
}

bool OAuth2ClientCredentialsAuthenticator::validate_token(const std::string& access_token) {
    return !access_token.empty() && access_token.length() > 10;
}

std::string OAuth2ClientCredentialsAuthenticator::get_scheme() const {
    return "Bearer";
}

std::vector<AuthChallenge> OAuth2ClientCredentialsAuthenticator::get_challenges() const {
    return {AuthChallenge{
        .scheme = "Bearer",
        .realm = "OAuth2 Client Credentials",
        .parameters = {
            {"token_url", config_.token_url}
        }
    }};
}

bool OAuth2ClientCredentialsAuthenticator::supports_refresh() const {
    return false; // Client credentials tokens don't typically support refresh
}

std::expected<OAuth2TokenResponse, OAuth2Error> OAuth2ClientCredentialsAuthenticator::get_client_credentials_token() {
    std::string post_data =
        "grant_type=client_credentials" +
        "&client_id=" + curl_easy_escape(nullptr, config_.client_id.c_str(), 0) +
        "&client_secret=" + curl_easy_escape(nullptr, config_.client_secret.c_str(), 0);

    if (!config_.scopes.empty()) {
        std::string scope_str;
        for (const auto& scope : config_.scopes) {
            if (!scope_str.empty()) scope_str += " ";
            scope_str += scope;
        }
        post_data += "&scope=" + curl_easy_escape(nullptr, scope_str.c_str(), 0);
    }

    auto response = http_request(config_.token_url, "POST", {{"Content-Type", "application/x-www-form-urlencoded"}}, post_data);

    if (!response || response->status_code != 200) {
        OAuth2Error error;
        error.error = "invalid_client";
        error.error_description = "Failed to obtain client credentials token";
        return error;
    }

    auto params = parse_form_encoded(response->body);
    OAuth2TokenResponse token_response;
    token_response.access_token = params["access_token"];
    token_response.token_type = params["token_type"];
    token_response.expires_in = std::stoi(params["expires_in"]);

    if (params.count("scope")) {
        token_response.scope = params["scope"];
    }

    return token_response;
}

// OAuth2DeviceAuthenticator implementation
OAuth2DeviceAuthenticator::OAuth2DeviceAuthenticator(const OAuth2Config& config)
    : config_(config) {}

Result<AuthResult> OAuth2DeviceAuthenticator::authenticate(const Credentials& credentials) {
    auto device_code_result = request_device_code();
    if (!device_code_result) {
        return Error::authentication_error("Failed to request device code: " +
            device_code_result.error().error_description.value_or("unknown error"));
    }

    auto device_code = device_code_result->device_code;

    // Poll for token
    auto start_time = std::chrono::steady_clock::now();
    auto timeout = std::chrono::seconds(device_code_result->expires_in);

    while (std::chrono::steady_clock::now() - start_time < timeout) {
        auto token_result = poll_for_token(device_code);
        if (token_result) {
            auto& token = token_result.value();
            AuthResult result;
            result.success = true;
            result.access_token = token.access_token;
            result.token_type = token.token_type;
            result.expires_in = token.expires_in;

            if (token.refresh_token) {
                result.refresh_token = *token.refresh_token;
            }

            return result;
        }

        auto& error = token_result.error();
        if (error.error != "authorization_pending") {
            return Error::authentication_error("Device authentication failed: " +
                error.error_description.value_or(error.error));
        }

        std::this_thread::sleep_for(std::chrono::seconds(device_code_result->interval));
    }

    return Error::timeout("Device authentication timed out");
}

Result<AuthResult> OAuth2DeviceAuthenticator::refresh(const std::string& refresh_token) {
    if (refresh_token.empty()) {
        return Error::validation_error("Refresh token is required");
    }

    std::string post_data =
        "grant_type=refresh_token" +
        "&refresh_token=" + curl_easy_escape(nullptr, refresh_token.c_str(), 0) +
        "&client_id=" + curl_easy_escape(nullptr, config_.client_id.c_str(), 0);

    if (!config_.client_secret.empty()) {
        post_data += "&client_secret=" + curl_easy_escape(nullptr, config_.client_secret.c_str(), 0);
    }

    auto response = http_request(config_.token_url, "POST", {{"Content-Type", "application/x-www-form-urlencoded"}}, post_data);

    if (!response || response->status_code != 200) {
        return Error::authentication_error("Failed to refresh device token");
    }

    auto params = parse_form_encoded(response->body);

    AuthResult result;
    result.success = true;
    result.access_token = params["access_token"];
    result.token_type = params["token_type"];
    result.expires_in = std::stoi(params["expires_in"]);

    if (params.count("refresh_token")) {
        result.refresh_token = params["refresh_token"];
    }

    return result;
}

bool OAuth2DeviceAuthenticator::validate_token(const std::string& access_token) {
    return !access_token.empty() && access_token.length() > 10;
}

std::string OAuth2DeviceAuthenticator::get_scheme() const {
    return "Bearer";
}

std::vector<AuthChallenge> OAuth2DeviceAuthenticator::get_challenges() const {
    return {AuthChallenge{
        .scheme = "Bearer",
        .realm = "OAuth2 Device Flow",
        .parameters = {
            {"token_url", config_.token_url}
        }
    }};
}

bool OAuth2DeviceAuthenticator::supports_refresh() const {
    return true;
}

std::expected<OAuth2DeviceAuthenticator::DeviceCodeResponse, OAuth2Error> OAuth2DeviceAuthenticator::request_device_code() {
    std::string device_url = config_.token_url;
    // Replace /token with /device/code in the URL
    size_t token_pos = device_url.find("/token");
    if (token_pos != std::string::npos) {
        device_url.replace(token_pos, 6, "/device/code");
    } else {
        device_url += "/device/code";
    }

    std::string post_data =
        "client_id=" + curl_easy_escape(nullptr, config_.client_id.c_str(), 0);

    if (!config_.scopes.empty()) {
        std::string scope_str;
        for (const auto& scope : config_.scopes) {
            if (!scope_str.empty()) scope_str += " ";
            scope_str += scope;
        }
        post_data += "&scope=" + curl_easy_escape(nullptr, scope_str.c_str(), 0);
    }

    auto response = http_request(device_url, "POST", {{"Content-Type", "application/x-www-form-urlencoded"}}, post_data);

    if (!response || response->status_code != 200) {
        OAuth2Error error;
        error.error = "invalid_request";
        error.error_description = "Failed to request device code";
        return error;
    }

    auto params = parse_form_encoded(response->body);
    DeviceCodeResponse device_response;
    device_response.device_code = params["device_code"];
    device_response.user_code = params["user_code"];
    device_response.verification_uri = params["verification_uri"];

    if (params.count("verification_uri_complete")) {
        device_response.verification_uri_complete = params["verification_uri_complete"];
    }

    if (params.count("expires_in")) {
        device_response.expires_in = std::stoi(params["expires_in"]);
    }

    if (params.count("interval")) {
        device_response.interval = std::stoi(params["interval"]);
    }

    return device_response;
}

std::expected<OAuth2TokenResponse, OAuth2Error> OAuth2DeviceAuthenticator::poll_for_token(
    const std::string& device_code) {

    std::string post_data =
        "grant_type=urn:ietf:params:oauth:grant-type:device_code" +
        "&device_code=" + curl_easy_escape(nullptr, device_code.c_str(), 0) +
        "&client_id=" + curl_easy_escape(nullptr, config_.client_id.c_str(), 0);

    if (!config_.client_secret.empty()) {
        post_data += "&client_secret=" + curl_easy_escape(nullptr, config_.client_secret.c_str(), 0);
    }

    auto response = http_request(config_.token_url, "POST", {{"Content-Type", "application/x-www-form-urlencoded"}}, post_data);

    if (!response) {
        OAuth2Error error;
        error.error = "network_error";
        error.error_description = "Failed to connect to token endpoint";
        return error;
    }

    if (response->status_code != 200) {
        auto params = parse_form_encoded(response->body);
        OAuth2Error error;
        error.error = params.count("error") ? params["error"] : "unknown";
        if (params.count("error_description")) {
            error.error_description = params["error_description"];
        }
        return error;
    }

    auto params = parse_form_encoded(response->body);
    OAuth2TokenResponse token_response;
    token_response.access_token = params["access_token"];
    token_response.token_type = params["token_type"];
    token_response.expires_in = std::stoi(params["expires_in"]);

    if (params.count("refresh_token")) {
        token_response.refresh_token = params["refresh_token"];
    }
    if (params.count("scope")) {
        token_response.scope = params["scope"];
    }

    return token_response;
}

} // namespace auth
} // namespace mcp